﻿using System;
using System.Collections.Concurrent;
using System.Collections.Generic;
using System.Linq;
using System.Text;
using System.Threading;
using System.Threading.Tasks;

namespace System.Threading.Tasks
{
    /// <summary>Provides a pump that supports running asynchronous methods on the current thread.</summary>
    /// <example><code><![CDATA[
    /// AsyncPump.Run(async delegate
    /// {
    ///     await DemoAsync();
    /// });
    /// 
    /// static async Task DemoAsync()
    /// {
    ///     var d = new Dictionary<int, int>();
    ///
    ///     for (int i = 0; i < 10000; i++)
    ///     {
    ///         int id = Thread.CurrentThread.ManagedThreadId;
    ///         int count;
    ///         d[id] = d.TryGetValue(id, out count) ? count + 1 : 1;
    ///         await Task.Yield();
    ///     }
    /// 
    ///     foreach (var pair in d)
    ///         Console.WriteLine(pair);
    /// }
    /// ]]>
    /// </code></example>
    public static class AsyncPump
    {
        /// <summary>Runs the specified asynchronous function.</summary>
        /// <param name="func">The asynchronous function to execute.</param>
        public static void Run(Func<Task> func)
        {
            if (func == null)
                throw new ArgumentNullException(nameof(func));

            var prevCtx = SynchronizationContext.Current;
            SingleThreadSynchronizationContext syncCtx = null;

            try
            {
                // Establish the new context
                syncCtx = new SingleThreadSynchronizationContext();
                SynchronizationContext.SetSynchronizationContext(syncCtx);

                // Invoke the function and alert the context to when it completes
                var t = func();
                if (t == null) throw new InvalidOperationException("No task provided.");
                t.ContinueWith(delegate { syncCtx.Complete(); }, TaskScheduler.Default);

                // Pump continuations and propagate any exceptions
                syncCtx.RunOnCurrentThread();
                t.GetAwaiter().GetResult();
            }
            finally
            {
                SynchronizationContext.SetSynchronizationContext(prevCtx);
                syncCtx?.Dispose();
            }
        }

        /// <summary>Provides a SynchronizationContext that's single-threaded.</summary>
        private sealed class SingleThreadSynchronizationContext : SynchronizationContext, IDisposable
        {
            /// <summary>The queue of work items.</summary>
            private readonly BlockingCollection<KeyValuePair<SendOrPostCallback, object>> _queue =
                new BlockingCollection<KeyValuePair<SendOrPostCallback, object>>();
            
            /// <summary>The processing thread.</summary>
            private readonly Thread _thread = Thread.CurrentThread;

            /// <summary>Dispatches an asynchronous message to the synchronization context.</summary>
            /// <param name="d">The System.Threading.SendOrPostCallback delegate to call.</param>
            /// <param name="state">The object passed to the delegate.</param>
            public override void Post(SendOrPostCallback d, object state)
            {
                if (d == null)
                    throw new ArgumentNullException("d");

                _queue.Add(new KeyValuePair<SendOrPostCallback, object>(d, state));
            }

            /// <summary>Not supported.</summary>
            public override void Send(SendOrPostCallback d, object state)
            {
                throw new NotSupportedException("Synchronously sending is not supported.");
            }

            /// <summary>Runs an loop to process all queued work items.</summary>
            public void RunOnCurrentThread()
            {
                foreach (var workItem in _queue.GetConsumingEnumerable())
                    workItem.Key(workItem.Value);
            }

            /// <summary>Notifies the context that no more work will arrive.</summary>
            public void Complete() { _queue.CompleteAdding(); }

            #region Dispose Pattern
            private bool _disposed;

            public void Dispose()
            {
                Dispose(true);
                GC.SuppressFinalize(this);
            }

            private void Dispose(bool disposing)
            {
                if (!_disposed)
                {
                    // if this is a dispose call dispose on all state you 
                    // hold, and take yourself off the Finalization queue.
                    if (disposing)
                    {
                        _queue.Dispose();
                    }

                    _disposed = true;
                }
            }
            #endregion


        }
    }
}
